{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25cde456",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "37826784",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import random\n",
    "import numpy as np\n",
    "\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from data_process.sort_process import SortProcess\n",
    "from data_process.data_iter import DataIter \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "\n",
    "\n",
    "SEED = 1\n",
    "random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "\n",
    "torch.backends.cudnn.deterministic = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12979490",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "id": "c9c68605",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Encoder(nn.Module):\n",
    "    def __init__(self, input_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout):\n",
    "        super().__init__()\n",
    "        self.embedding = nn.Embedding(input_dim, emb_dim)\n",
    "        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)\n",
    "        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, src): \n",
    "        '''\n",
    "        src = [src_len, batch_size]\n",
    "        '''\n",
    "        src = src.transpose(0, 1) # src = [batch_size, src_len]\n",
    "        embedded = self.dropout(self.embedding(src)).transpose(0, 1) # embedded = [src_len, batch_size, emb_dim]\n",
    "        \n",
    "        # enc_output = [src_len, batch_size, hid_dim * num_directions]\n",
    "        # enc_hidden = [n_layers * num_directions, batch_size, hid_dim]\n",
    "        enc_output, enc_hidden = self.rnn(embedded) # if h_0 is not give, it will be set 0 acquiescently\n",
    "\n",
    "        # enc_hidden is stacked [forward_1, backward_1, forward_2, backward_2, ...]\n",
    "        # enc_output are always from the last layer\n",
    "        \n",
    "        # enc_hidden [-2, :, : ] is the last of the forwards RNN \n",
    "        # enc_hidden [-1, :, : ] is the last of the backwards RNN\n",
    "        \n",
    "        # initial decoder hidden is final hidden state of the forwards and backwards \n",
    "        # encoder RNNs fed through a linear layer\n",
    "        # s = [batch_size, dec_hid_dim]\n",
    "        s = torch.tanh(self.fc(torch.cat((enc_hidden[-2,:,:], enc_hidden[-1,:,:]), dim = 1)))\n",
    "        \n",
    "        return enc_output, s \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "4a745466",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self, enc_hid_dim, dec_hid_dim):\n",
    "        super().__init__()\n",
    "        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim, bias=False)\n",
    "        self.v = nn.Linear(dec_hid_dim, 1, bias = False)\n",
    "        \n",
    "    def forward(self, s, enc_output):\n",
    "        \n",
    "        # s = [batch_size, dec_hid_dim]\n",
    "        # enc_output = [src_len, batch_size, enc_hid_dim * 2]\n",
    "        \n",
    "        batch_size = enc_output.shape[1]\n",
    "        src_len = enc_output.shape[0]\n",
    "        \n",
    "        # repeat decoder hidden state src_len times\n",
    "        # s = [batch_size, src_len, dec_hid_dim]\n",
    "        # enc_output = [batch_size, src_len, enc_hid_dim * 2]\n",
    "        s = s.unsqueeze(1).repeat(1, src_len, 1)\n",
    "        enc_output = enc_output.transpose(0, 1)\n",
    "        \n",
    "        # energy = [batch_size, src_len, dec_hid_dim]\n",
    "        energy = torch.tanh(self.attn(torch.cat((s, enc_output), dim = 2)))\n",
    "        \n",
    "        # attention = [batch_size, src_len]\n",
    "        attention = self.v(energy).squeeze(2)\n",
    "        \n",
    "        return F.softmax(attention, dim=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "id": "78381e46",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Seq2Seq(nn.Module):\n",
    "    def __init__(self, encoder, decoder, device):\n",
    "        super().__init__()\n",
    "        self.encoder = encoder\n",
    "        self.decoder = decoder\n",
    "        self.device = device\n",
    "        \n",
    "    def forward(self, src, trg, teacher_forcing_ratio = 0.5):\n",
    "        \n",
    "        # src = [src_len, batch_size]\n",
    "        # trg = [trg_len, batch_size]\n",
    "        # teacher_forcing_ratio is probability to use teacher forcing\n",
    "        \n",
    "        batch_size = src.shape[1]\n",
    "        trg_len = trg.shape[0]\n",
    "        trg_vocab_size = self.decoder.output_dim\n",
    "        \n",
    "        # tensor to store decoder outputs\n",
    "        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(self.device)\n",
    "        \n",
    "        # enc_output is all hidden states of the input sequence, back and forwards\n",
    "        # s is the final forward and backward hidden states, passed through a linear layer\n",
    "        enc_output, s = self.encoder(src)\n",
    "                \n",
    "        # first input to the decoder is the <sos> tokens\n",
    "        dec_input = trg[0,:]\n",
    "        \n",
    "        for t in range(1, trg_len):\n",
    "            \n",
    "            # insert dec_input token embedding, previous hidden state and all encoder hidden states\n",
    "            # receive output tensor (predictions) and new hidden state\n",
    "            dec_output, s = self.decoder(dec_input, s, enc_output)\n",
    "            \n",
    "            # place predictions in a tensor holding predictions for each token\n",
    "            outputs[t] = dec_output\n",
    "            \n",
    "            # decide if we are going to use teacher forcing or not\n",
    "            teacher_force = random.random() < teacher_forcing_ratio\n",
    "            \n",
    "            # get the highest predicted token from our predictions\n",
    "            top1 = dec_output.argmax(1) \n",
    "            \n",
    "            # if teacher forcing, use actual next token as next input\n",
    "            # if not, use predicted token\n",
    "            dec_input = trg[t] if teacher_force else top1\n",
    "\n",
    "        return outputs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16c2950b",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "07cecc8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Decoder(nn.Module):\n",
    "    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):\n",
    "        super().__init__()\n",
    "        self.output_dim = output_dim\n",
    "        self.attention = attention\n",
    "        self.embedding = nn.Embedding(output_dim, emb_dim)\n",
    "        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)\n",
    "        self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)\n",
    "        self.dropout = nn.Dropout(dropout)\n",
    "        \n",
    "    def forward(self, dec_input, s, enc_output):\n",
    "             \n",
    "        # dec_input = [batch_size]\n",
    "        # s = [batch_size, dec_hid_dim]\n",
    "        # enc_output = [src_len, batch_size, enc_hid_dim * 2]\n",
    "        \n",
    "        dec_input = dec_input.unsqueeze(1) # dec_input = [batch_size, 1]\n",
    "        \n",
    "        embedded = self.dropout(self.embedding(dec_input)).transpose(0, 1) # embedded = [1, batch_size, emb_dim]\n",
    "        \n",
    "        # a = [batch_size, 1, src_len]  \n",
    "        a = self.attention(s, enc_output).unsqueeze(1)\n",
    "        \n",
    "        # enc_output = [batch_size, src_len, enc_hid_dim * 2]\n",
    "        enc_output = enc_output.transpose(0, 1)\n",
    "\n",
    "        # c = [1, batch_size, enc_hid_dim * 2]\n",
    "        c = torch.bmm(a, enc_output).transpose(0, 1)\n",
    "\n",
    "        # rnn_input = [1, batch_size, (enc_hid_dim * 2) + emb_dim]\n",
    "        rnn_input = torch.cat((embedded, c), dim = 2)\n",
    "            \n",
    "        # dec_output = [src_len(=1), batch_size, dec_hid_dim]\n",
    "        # dec_hidden = [n_layers * num_directions, batch_size, dec_hid_dim]\n",
    "        dec_output, dec_hidden = self.rnn(rnn_input, s.unsqueeze(0))\n",
    "        \n",
    "        # embedded = [batch_size, emb_dim]\n",
    "        # dec_output = [batch_size, dec_hid_dim]\n",
    "        # c = [batch_size, enc_hid_dim * 2]\n",
    "        embedded = embedded.squeeze(0)\n",
    "        dec_output = dec_output.squeeze(0)\n",
    "        c = c.squeeze(0)\n",
    "        \n",
    "        # pred = [batch_size, output_dim]\n",
    "        pred = self.fc_out(torch.cat((dec_output, c, embedded), dim = 1))\n",
    "        \n",
    "        return pred, dec_hidden.squeeze(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "435e44c2",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
